
import os
import time
import shutil
from pypinyin import pinyin, Style, lazy_pinyin
from run_uvr import *
from run_infer import *
from ModelConfig import *
from song_constant import get_song_resources


def time_cost(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()  # 记录函数开始执行的时间
        result = func(*args, **kwargs)  # 执行函数
        end_time = time.time()  # 记录函数执行结束的时间
        elapsed_time = end_time - start_time  # 计算耗时
        print(f"函数执行 {func.__name__} 时间为 {elapsed_time} seconds  .")
        return result
    return wrapper
def build_tmp_folder(song_path):
    # 设置文件夹名称和路径
    folder_name = "bksong"

    # 检查文件夹是否存在
    if os.path.exists(folder_name):
        # 如果存在，则删除该文件夹
        print(f"The folder {folder_name} already exists. Deleting it...")
        shutil.rmtree(folder_name)
        print(f"The folder {folder_name} has been deleted.")

    # 创建一个新的空文件夹
    print(f"Creating a new folder {folder_name}...")
    os.makedirs(folder_name)
    print(f"The folder {folder_name} has been created.")

    # 设置源文件路径和目标文件夹路径
    source_file = song_path  # 'path/to/your/source/file.txt'  # 替换为你的源文件路径
    destination_folder = folder_name  # 'path/to/your/destination/folder'  # 替换为你的目标文件夹路径

    # 确保目标文件夹存在
    if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)
    # 分割文件路径，获取文件名
    file_name = os.path.basename(song_path)

    # 移除文件扩展名，获取不带扩展名的文件名
    file_name_without_extension = os.path.splitext(file_name)[0]
    pinyin_list = lazy_pinyin(file_name_without_extension)
    print(pinyin_list, os.path.splitext(file_name)[1])
    pinyin_str = (''.join(pinyin_list) + os.path.splitext(file_name)[1]).replace(' ', '').replace('-', '')
    print(pinyin_str)
    print('pinyin_str', pinyin_str)
    # 构建目标文件的完整路径
    source_file = pinyin_str

    destination_file = os.path.join(destination_folder, os.path.basename(source_file))

    # 复制文件
    print({'song': song_path, 'copyTo': destination_file})
    shutil.copy2(song_path, destination_file)

    print(f"The file has been copied to {destination_file}")
    return destination_folder, pinyin_str

def get_split_model_outputs(save_root_vocal,save_root_ins,file_name_without_extension):
    """
    根据output 文件夹取出对应文件
    :param save_root_vocal:
    :param save_root_ins:
    :param file_name_without_extension:
    :return:
    """
    vocal_files = [e for e in os.listdir(save_root_vocal) if file_name_without_extension in e and '.wav' in e[-5:]]
    instrc_files = [e for e in os.listdir(save_root_ins) if file_name_without_extension in e and '.wav' in e[-5:]]
    print('vocal_files', vocal_files)
    print('instrc_files', instrc_files)
    instrc_file = os.path.join(save_root_vocal, vocal_files[0])
    vocal_file = os.path.join(save_root_ins, instrc_files[0])
    return instrc_file,vocal_file

def run_task_infer(task_name,config):
    task = VoiceConversionTask(task=task_name)
    print(f'Model path: {task.model_path}')
    print(f'Index file path: {task.index_file_path}')
    return run_infer_model(task.model_path,
                           task.index_file_path,
                           config.raw_instrument_file,
                           config.raw_vocal_file,
                           config.vocal_ai_gen_path  ,
                           config.song_gen_path)

@time_cost
def call_generation(task_id,user_id,song_id,config):

    instrc_file, vocal_file=  get_song_resources(song_id)
    config.raw_instrument_file = instrc_file
    config.raw_vocal_file = vocal_file
    print('resource ',instrc_file,vocal_file)
    run_task_infer(task_id, config)
    print('done generation',os.path.exists(os.path.abspath(config.song_gen_path)))
    return config


def run_sample2(model_trained_task, song_path, output='huazhou_output.wav'):
    model_name = 'HP3_all_vocals'
    inp_root, filename = build_tmp_folder(song_path)

    file_name_without_extension = os.path.splitext(filename)[0]
    format0 = filename.split('.')[-1]

    split_outs = AudioSplitPaths(file_name_without_extension)
    print(f'Save root vocal path: {split_outs.save_root_vocal}')
    print(f'Save root ins path: {split_outs.save_root_ins}')
    agg = 10

    print('run input folder', inp_root, format0)
    print(inp_root, split_outs.save_root_vocal, split_outs.save_root_ins, format0)
    print({'inp': inp_root, 'vocal_output': split_outs.save_root_vocal,
           'instru_output': split_outs.save_root_ins, 'format0': format0

           })
    split_outs.make_folder()
    # autodl-tmp/Retrieval-based-Voice-Conversion-WebUI/tmp_song
    run_model(os.path.abspath(inp_root),
              split_outs.save_root_vocal,
              split_outs.save_root_ins,
              format0)

    instrc_file,vocal_file  = get_split_model_outputs(split_outs.save_root_vocal,
                                                      split_outs.save_root_ins,file_name_without_extension)
    print('instr', instrc_file, 'vocal', vocal_file)

    run_task_infer(model_trained_task, instrc_file, vocal_file, output)

# run_sample2('/root/autodl-tmp/infer_song/谢春花-我从崖边跌落.mp3', '我从崖边跌落.wav')
# run_sample2('/root/autodl-tmp/sample_train/回龙大街.m4a')
